import gradio as gr
from PIL import Image
import requests
import time

import io
from results import results
from util.utils import get_AccessKey_SecretKey, upload_image,make_sign
from copy import deepcopy
from pathlib import Path
from step import Step

AccessKey, SecretKey = get_AccessKey_SecretKey()

class Step3(Step):
    # 预设风格到 modelId 的映射（用于前端下拉框）
    LORA_STYLES = {
        "卡通": "5aad2800df224473acbd27d92aea3f3f",
        "插画": "7aa06b226feb46f485a6793a8d5a5184",
        "水墨": "2bc8ff1e8bc847008fd40e40efcdd096",
        "海报": "31360f2f031b4ff6b589412a52713fcf",
        "朋克": "0ad44fc3ca564bba864c82a36f3a8f65",
    }
    API_URL = "https://openapi.liblibai.cloud/api/generate/webui/img2img"
    STATUS_URL = "https://api.liblib.ai/api/generate/webui/status"

    PARAMS = {
    "templateUuid": "9c7d531dc75f476aa833b3d452b8f7ad",
    "generateParams": {
        "checkPointId": "bd065cff3a854af2b28659ed0f6d289d",
        "prompt": "1 girl wear sunglasses",
        "negativePrompt": "",
        "clipSkip": 2,
        "sampler": 15,
        "steps": 20,
        "cfgScale": 7,
        "randnSource": 0,
        "seed": -1,
        "imgCount": 1,
        "restoreFaces": 0,
        "sourceImage": "https://liblibai-online.liblib.cloud/img/081e9f07d9bd4c2ba090efde163518f9/7c1cc38e-522c-43fe-aca9-07d5420d743e.png",
        "resizeMode": 1,
        "resizedWidth": 1024,
        "resizedHeight": 1024,
        "mode": 0,
        "denoisingStrength": 0.75,
        "additionalNetwork": [
            {
                "modelId": "0ad44fc3ca564bba864c82a36f3a8f65",
                "weight": 0
            },
            {
                "modelId": "5aad2800df224473acbd27d92aea3f3f", 
                "weight": 0.6 
            }
             ],

    }
}
    def __init__(self):
        super().__init__("step3")
        
    # -----------------------------------------------------
    # ✅ 调用图生图接口
    # -----------------------------------------------------
    def send_img2img(self,payload):
        headers = {
            "Content-Type": "application/json",
        }
        sign,timestamp, signature_nonce=make_sign(SecretKey,type="img2img")
        url = self.API_URL + f"?AccessKey={AccessKey}&Signature={sign}&SignatureNonce={signature_nonce}&Timestamp={timestamp}"
        resp = requests.post(url, headers=headers, json=payload)
        data = resp.json()

        if data.get("code") != 0:
            print("请求失败:", data)
            raise RuntimeError(f"请求失败: {data}")
            

        return data["data"]["generateUuid"]


    # -----------------------------------------------------
    # ✅ 轮询查询结果
    # -----------------------------------------------------
    def query_result(self, uuid):
        headers = {
            "Content-Type": "application/json",
        }

        while True:
            resp = requests.post(self.STATUS_URL, headers=headers, json={"generateUuid": uuid})
            data = resp.json()

            status = data["data"]["generateStatus"]
            percent = data["data"]["percentCompleted"]

            if status == 5:
                images = data["data"]["images"]
                return [img["imageUrl"] for img in images]

            yield f"生成中... {percent}%"
            time.sleep(2)


    # -----------------------------------------------------
    # ✅ Gradio 点击事件：收集参数 → 调API → 返回生成图片
    # -----------------------------------------------------
    def run_generate(self,
        source_image:Image.Image,
        prompt: str,
        steps: int,
        cfg_scale: float,
        denoise: float,
        lora1_id: str = None,
        lora1_weight: float = 0.0,
        lora2_id: str = None,
        lora2_weight: float = 0.0
    ):
        global AccessKey,SecretKey
        AccessKey, SecretKey = get_AccessKey_SecretKey()
        if AccessKey is None or SecretKey is None:
            gr.Warning("请先在[设置]中创建并选择一个用户以存储密钥",visible=True,duration=5)
            return None
        
        if source_image is None:
            return "请上传一张源图"


        # ✅ 最终生图参数
        # 使用 deepcopy 避免修改全局 PARAMS
        payload = deepcopy(self.PARAMS)

        # 处理前端传入的 LoRA（前端传的是风格名或 "无"，需映射为 modelId）
        additional = []
        try:
            # 如果用户选择了预设风格，则从 LORA_STYLES 映射到 modelId；否则允许用户直接传入 modelId（向后兼容）
            if lora1_id and str(lora1_id).strip() not in ("", "无"):
                model_id_1 = self.LORA_STYLES.get(str(lora1_id).strip(), str(lora1_id).strip())
                if model_id_1:
                    additional.append({"modelId": model_id_1, "weight": float(lora1_weight or 0.0)})
        except Exception:
            pass
        try:
            if lora2_id and str(lora2_id).strip() not in ("", "无"):
                model_id_2 = self.LORA_STYLES.get(str(lora2_id).strip(), str(lora2_id).strip())
                if model_id_2:
                    additional.append({"modelId": model_id_2, "weight": float(lora2_weight or 0.0)})
        except Exception:
            pass

        if additional:
            payload.setdefault("generateParams", {})["additionalNetwork"] = additional

        payload["generateParams"]["sourceImage"] = upload_image(AccessKey,SecretKey,source_image)
        payload["generateParams"]["prompt"] = prompt
        payload["generateParams"]["steps"] = steps
        payload["generateParams"]["cfgScale"] = cfg_scale
        payload["generateParams"]["denoisingStrength"] = denoise

        # ✅ 发起生图任务
        generateUuid = self.send_img2img(payload)
        print("生成任务 UUID:", generateUuid)

        headers = {
            "Content-Type": "application/json",
        }
        sign, timestamp, signature_nonce = make_sign(SecretKey,type="status")
        get_url = "https://openapi.liblibai.cloud/api/generate/comfy/status"
        get_url = f"{get_url}?AccessKey={AccessKey}&Signature={sign}&SignatureNonce={signature_nonce}&Timestamp={timestamp}"
        data = {"generateUuid": generateUuid}
        while True:
            resp = requests.post(get_url, json=data, headers=headers)
            if resp.status_code == 200:
                res = resp.json()
                if res.get("code") != 0:
                    continue
                if res.get("data").get("generateStatus") == 5:
                    out = res.get("data").get("images")[0]['imageUrl']
                    print("generated img url:", out)
                    img_response = requests.get(out)
                    results.step3_result_image = Image.open(io.BytesIO(img_response.content))
                    self.save_img(results.step3_result_image, generateUuid)
                    return Image.open(io.BytesIO(img_response.content))
            time.sleep(1)

        return None

    def create_step(self):

        with gr.Blocks(title="LiblibAI Img2Img") as demo:

            gr.Markdown("## LiblibAI 图生图测试界面（图片上传版）")

            with gr.Row():
                source_image = gr.Image(type="pil", label="源图上传")
                result_image = gr.Image(type='pil',label="生成结果")

            with gr.Row():
                prompt = gr.Textbox(label="Prompt", value="1 girl wear sunglasses")

            with gr.Row():
                steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1)
                cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=15, value=7, step=0.5)
                denoise = gr.Slider(label="去噪强度（Denoising Strength）", minimum=0.0, maximum=1.0, value=0.75, step=0.01)
            
            gr.Markdown("最多可添加 5 个 LoRA（示例两组，可扩展）")

            with gr.Column():
                # 使用下拉框选择风格，前端选择一个风格名，后端映射为对应的 modelId
                lora1_id = gr.Dropdown(choices=["无", "卡通", "插画", "水墨", "海报", "朋克"], value="卡通", label="LoRA 1 风格")
                lora1_weight = gr.Slider(0, 1, value=0.6, label="LoRA 1 权重")

                lora2_id = gr.Dropdown(choices=["无", "卡通", "插画", "水墨", "海报", "朋克"], value="水墨", label="LoRA 2 风格")
                lora2_weight = gr.Slider(0, 1, value=0.3, label="LoRA 2 权重")

        
            btn = gr.Button("开始生成")
            load = gr.Button("加载上一步图片",variant="primary")

            btn.click(
                self.run_generate,
                inputs=[source_image, prompt, steps, cfg_scale, denoise, lora1_id, lora1_weight, lora2_id, lora2_weight],
                outputs=result_image
            )
            def _load():
                return results.step2_result_image

            load.click(
                _load,
                outputs=[source_image]
            )

            return demo

if __name__ == "__main__":
    step3 = Step3()
    step3.create_step().launch()
